package cn.doitedu.dwetl.utls

import java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream}
import java.nio.ByteBuffer

import cn.doitedu.commons.utils.SparkUtil
import org.apache.spark.rdd.RDD
import org.roaringbitmap.{RoaringBitmap, Util}

object RoaringBitMapSpark {

  def main(args: Array[String]): Unit = {
    val spark = SparkUtil.getSparkSession("")
    import spark.implicits._
    val df = spark.createDataset(Seq(
      ("123","a1","b1"),
      ("126","a1","b1"),
      ("233","a1","b2"),
      ("123","a1","b2"),
      ("123","a2","b1"),
      ("233","a2","b1"),
      ("154","a3","b1"),
      ("133","a3","b2"),
      ("154","a3","b2")
    )).toDF("guid","d1","d2")

    df.createTempView("df")
    spark.sql(
      """
        |select
        |d1,d2,count(distinct guid) as uv
        |from df
        |group by d1,d2
        |
        |
        |""".stripMargin).show(100,false)

    spark.sql(
      """
        |select
        |d1,count(distinct guid) as uv
        |from df
        |group by d1
        |
        |
        |""".stripMargin).show(100,false)

    val rd = df.rdd.map(row=>{
      val guid = row.getAs[String]("guid")
      val d1 = row.getAs[String]("d1")
      val d2 = row.getAs[String]("d2")
      (d1,d2,guid)
    })

    val rd2: RDD[((String, String), Array[Byte])] = rd.groupBy(tp=>(tp._1,tp._2)).mapValues(it=>{
      val bm = RoaringBitmap.bitmapOf(it.toArray.map(_._3.hashCode):_*)
      val bo = new ByteArrayOutputStream(bm.serializedSizeInBytes())
      val oo = new ObjectOutputStream(bo)
      oo.writeObject(bm)
      bo.toByteArray
    })

    rd2.cache()
    val res = rd2.map(tp=>(tp._1._1,tp._1._2,tp._2)).toDF("d1","d2","bm")
    res.printSchema()
    res.coalesce(1).write.parquet("testdata/rrmpout/")


    val rd3: RDD[(String, RoaringBitmap)] = rd2.groupBy(tp=>tp._1._1).mapValues(iter=>{
      val bms: List[RoaringBitmap] = iter.map(_._2).map(ba=>{
        val bi = new ByteArrayInputStream(ba)
        val oi = new ObjectInputStream(bi)
        val bm = oi.readObject().asInstanceOf[RoaringBitmap]
        bm
      }).toList

      val bm0 = bms(0)
      for(i <- 1 until bms.size)
      bm0.or(bms(i))
      bm0
    })

    rd3.map(tp=>{
      (tp._1,tp._2.getCardinality)
    }).toDF("d1","uv").show(100,false)



    spark.close()
  }


}
